# masked_input
import random
import time

import torch

from pypbc import *
import numpy as np
import hashlib
from concurrent.futures import ProcessPoolExecutor, as_completed
from multiprocessing import cpu_count
Hash = hashlib.sha256

def Flatten_paremeter(w):
    flattened_params = []
    for param_tensor in w.values():
        flattened_params.extend(param_tensor.view(-1))
    flattened_params_tensor = torch.tensor(flattened_params)
    quantify_flattened_params_tensor = (flattened_params_tensor * 10000).int()
    quantify_flattened_params_list = quantify_flattened_params_tensor.numpy().tolist()
    return quantify_flattened_params_list

def KeyGen(qbits, rbits):
    params = Parameters(qbits=qbits, rbits=rbits)
    pairing = Pairing(params)
    g2 = Element.random(pairing, G2)
    return [G2, g2, params, pairing]


def Hhash(pp, weight, len_weight, pairing):
    hhash_list = []
    for i in range(len_weight):
        if weight[i]<0:
            Hhash_i = Element(pairing, G2, value = pp[i] ** abs(weight[i]))
            Hhash_i = Element.__invert__(Hhash_i)
        else:
            Hhash_i = Element(pairing, G2, value=pp[i] ** weight[i])
        hhash_list.append(Hhash_i)
    product = hhash_list[len_weight - 1]
    for i in range(len_weight - 1):
        product = Element.__mul__(product, hhash_list[i])
    return product

def PPGen(g2, params, pairing, dimen):
    pp = []
    for number in range(dimen):
        r = Element.random(pairing, Zr)
        g2_i = Element(pairing, G2, value=g2 ** r)
        pp.append(g2_i)
    return pp


def generate_g2_element(params, g2, index):
    params = Parameters(params)
    pairing=Pairing(params)
    g2 = Element(pairing, G2, g2)
    r = Element.random(pairing, Zr)
    g2_i = Element(pairing, G2, value=g2 ** r)
    g2_i = Element.__str__(g2_i)

    # g2_i = "1"
    return index, g2_i


def PPGen1(g2, params, pairing, dimen):
    pp = [None] * dimen  # 初始化一个列表来存放结果

    num_cores = cpu_count()  # 获取 CPU 核心数
    # print(num_cores)
    with ProcessPoolExecutor(max_workers=num_cores) as executor:

        # 提交所有任务
        future_to_index = {executor.submit(generate_g2_element, str(params), Element.__str__(g2), i): i for i in range(dimen)}

        for future in as_completed(future_to_index):
            index = future_to_index[future]
            try:
                # step+=1
                result_index, g2_i = future.result()
                pp[result_index] = g2_i
                # print(step)
            except Exception as exc:
                print(f'生成第 {index} 个元素时发生错误: {exc}')

    return pp

def sign_key_gen(G2, g2, pairing):
    sk_sign_1 = Element.random(pairing, Zr)
    pk_sign_1 = Element(pairing, G2, value=g2 ** sk_sign_1)
    return sk_sign_1, pk_sign_1

def sign(sk_user, pairing, h2_user):
    sig1_user = Element(pairing, G2, value = h2_user ** sk_user)
    return sig1_user

def Verify_agg_sign_1(pairing, h2_idxs, sig1_s, pk_sign_1_s, num_user, g2):
    pk_sign_1s = []
    for i in range(len(pk_sign_1_s)):
        if pk_sign_1_s[i] != 0 :
            pk_sign_1s.append(pk_sign_1_s[i])
    e2_list = []
    for i in range(num_user):
        e2_list.append(pairing.apply(pk_sign_1s[i], h2_idxs[i]))
    e2 = e2_list[num_user - 1]
    agg_sign_1 = sig1_s[num_user - 1]
    for i in range(num_user - 1):
        e2 = Element.__mul__(e2, e2_list[i])
        agg_sign_1 = Element.__mul__(agg_sign_1, sig1_s[i])

    e1 = pairing.apply(g2, agg_sign_1)
    if e1 == e2:
        return 1
    elif e1 != e2:
        return 0

def Verify_agg_sign_2(enc_w_new_s_recover, pairing, sig1_s_recover, pk_sign_2_s, num_user, g2, sig2_s_recover):
    pk_sign_2_s_recover = []
    for pks in pk_sign_2_s:
        if pks != 0:
            pk_sign_2_s_recover.append(pks)
    e2_list, e1_list = [], []
    for i in range(num_user):
        local_BFV_str = str(enc_w_new_s_recover[i])
        ver_str_enc_w_new_sigm_idx = local_BFV_str + str(sig1_s_recover[i])
        h3_idx_str = Hash(ver_str_enc_w_new_sigm_idx.encode()).hexdigest()
        h3_idx = Element.from_hash(pairing, G2, h3_idx_str)
        e2_list.append(pairing.apply(pk_sign_2_s_recover[i], h3_idx))

    e2 = e2_list[num_user - 1]
    agg_sign_2 = sig2_s_recover[num_user - 1]
    for i in range(num_user - 1):
        e2 = Element.__mul__(e2, e2_list[i])
        agg_sign_2 = Element.__mul__(agg_sign_2, sig2_s_recover[i])
    e1 = pairing.apply(g2, agg_sign_2)
    if e1 == e2:
        return 1
    elif e1 != e2:
        return 0

def Verify_sign_1(pairing, g2, sig1, pk_sign, h2u):
    e1 = pairing.apply(g2, sig1)
    e2 = pairing.apply(pk_sign, h2u)
    if e1 == e2:
        return 1
    elif e1 != e2:
        return 0

def decompose_agg_w_new(decrypted_agg_w_new, len_h1_list, all_num_users):
    decrypted_agg_w = decrypted_agg_w_new[len_h1_list:]
    h1_idx_list = []
    for idx in range(all_num_users):
        h1_idx_ASCII_list = decrypted_agg_w_new[idx*130:(idx+1)*130]
        h1_idx_chr_list = [chr(mess) for mess in h1_idx_ASCII_list]
        h1_idx_chr_str = ''.join('%s' % mess for mess in h1_idx_chr_list)
        h1_idx_list.append(h1_idx_chr_str)
    return decrypted_agg_w, h1_idx_list

def decompose_trac_w_new(trac_decr_w_new, idx_1, len_h1_list, all_num_users):
    len_hash = 130
    decrypted_w_loc = trac_decr_w_new[len_hash * all_num_users:]
    idx_1 = int(idx_1)
    h1_idx_ASCII_list = trac_decr_w_new[idx_1 * len_hash:(idx_1 + 1) * len_hash]
    h1_idx_chr_list = [chr(mess) for mess in h1_idx_ASCII_list]
    h1_idx_chr_str = ''.join('%s' % mess for mess in h1_idx_chr_list)
    return decrypted_w_loc, h1_idx_chr_str, h1_idx_ASCII_list

def Hhash_equa(pp1, decrypted_agg_w, len_decrypted_agg_w, pairing, G2, h1_idx_list_recover):
    h_1_1 = Hhash(pp1, decrypted_agg_w, len_decrypted_agg_w, pairing)

    h1_recover = Element(pairing, G2, value=h1_idx_list_recover[-1])
    for j in range(len(h1_idx_list_recover) - 1):
        h1_idx_recover_temp = Element(pairing, G2, value=h1_idx_list_recover[j])
        h1_recover = Element.__mul__(h1_recover, h1_idx_recover_temp)
    if h_1_1 == h1_recover:
        return 1
    else:
        return 0

def Invert_Quality(decrypted_agg_w, num_clients):
    len_weigh = len(decrypted_agg_w)
    for i in range(len_weigh):
        decrypted_agg_w[i] = decrypted_agg_w[i] / 10000 / num_clients
    return decrypted_agg_w

def Restore_values_to_w(w, quantify_flattened_params_list):
    index = 0

    for key in w.keys():
        param_tensor = w[key]
        shape = param_tensor.shape
        new_param_tensor = np.array(quantify_flattened_params_list[index:index + param_tensor.numel()]).reshape(shape)
        w[key] = torch.tensor(new_param_tensor)
        index += param_tensor.numel()
    return w
